# -*- coding: utf-8 -*-
"""
验证X-Date头问题
检查volcengine SDK是否正确添加X-Date头
"""

import os
import sys
import json
from pathlib import Path
from dotenv import load_dotenv
from datetime import datetime
from email.utils import formatdate

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 加载环境变量
load_dotenv(project_root / '.env')

# 导入必要的模块
from volcengine.auth.SignerV4 import SignerV4
from volcengine.auth.SignParam import SignParam
from volcengine.Credentials import Credentials
import requests
from collections import OrderedDict


def test_sign_with_xdate():
    """测试签名是否包含X-Date头"""
    print("=" * 80)
    print("测试 volcengine SDK 签名是否添加 X-Date 头")
    print("=" * 80)

    # 从环境变量读取配置
    access_key = os.getenv('HIAGENT_ACCESS_KEY')
    secret_key = os.getenv('HIAGENT_SECRET_KEY')
    host = os.getenv('HIAGENT_HOST', 'open.volcengineapi.com')
    account_id = os.getenv('HIAGENT_ACCOUNT_ID')
    service = 'llmops'
    region = 'cn-beijing'
    version = '2024-05-01'

    if not all([access_key, secret_key, account_id]):
        print("\n[ERROR] 缺少必要的环境变量:")
        print(f"  HIAGENT_ACCESS_KEY: {'已设置' if access_key else '未设置'}")
        print(f"  HIAGENT_SECRET_KEY: {'已设置' if secret_key else '未设置'}")
        print(f"  HIAGENT_ACCOUNT_ID: {'已设置' if account_id else '未设置'}")
        return False

    print(f"\n[INFO] 配置信息:")
    print(f"  Host: {host}")
    print(f"  Service: {service}")
    print(f"  Region: {region}")
    print(f"  Version: {version}")
    print(f"  Account ID: {account_id}")

    # 创建签名参数
    param = SignParam()
    param.path = "/"
    param.method = "POST"
    param.host = host

    # 设置查询参数
    query = OrderedDict()
    query['Action'] = 'ListAppCenterCategory'
    query['Version'] = version
    query['X-Account-Id'] = account_id
    param.query = query

    # 设置请求头（不包含X-Date，让SDK自动添加）
    headers = OrderedDict()
    headers['Host'] = host
    headers['Content-Type'] = 'application/json'
    param.header_list = headers
    param.headers = headers

    # 设置请求体
    body = {"Top": {}}
    param.body = json.dumps(body)

    print(f"\n[INFO] 签名前的请求头:")
    for key, value in param.headers.items():
        print(f"  {key}: {value}")

    # 执行签名
    print(f"\n[INFO] 执行签名...")
    signer = SignerV4()
    credentials = Credentials(access_key, secret_key, service, region)

    try:
        signer.sign(param, credentials)
        print(f"[SUCCESS] 签名成功")
    except Exception as e:
        print(f"[ERROR] 签名失败: {e}")
        import traceback
        traceback.print_exc()
        return False

    # 检查签名后的请求头
    print(f"\n[INFO] 签名后的请求头:")
    has_xdate = False
    for key, value in param.headers.items():
        print(f"  {key}: {value}")
        if key.lower() == 'x-date':
            has_xdate = True

    if not has_xdate:
        print(f"\n[WARNING] 签名后的请求头中没有 X-Date!")
        print(f"[INFO] 尝试手动添加 X-Date 头...")

        # 手动添加X-Date头
        param.headers['X-Date'] = formatdate(timeval=None, localtime=False, usegmt=True)
        param.header_list['X-Date'] = param.headers['X-Date']

        print(f"  X-Date: {param.headers['X-Date']}")

        # 重新签名
        print(f"\n[INFO] 重新签名...")
        signer.sign(param, credentials)
        print(f"[SUCCESS] 重新签名成功")

        print(f"\n[INFO] 重新签名后的请求头:")
        for key, value in param.headers.items():
            print(f"  {key}: {value}")

    # 尝试发送请求
    print(f"\n[INFO] 发送测试请求...")
    url = f"https://{host}/?Action=ListAppCenterCategory&Version={version}&X-Account-Id={account_id}"

    try:
        response = requests.post(
            url,
            headers=dict(param.headers),
            data=param.body,
            timeout=30,
            verify=False
        )

        print(f"[INFO] 响应状态码: {response.status_code}")
        print(f"\n[INFO] 响应内容:")
        try:
            result = response.json()
            print(json.dumps(result, indent=2, ensure_ascii=False))

            if response.status_code == 200:
                print(f"\n[SUCCESS] 请求成功!")
                return True
            else:
                print(f"\n[ERROR] 请求失败!")
                return False

        except json.JSONDecodeError:
            print(response.text)
            return False

    except Exception as e:
        print(f"[ERROR] 请求失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_manual_xdate():
    """测试手动添加X-Date头的方式"""
    print("\n" + "=" * 80)
    print("测试手动添加 X-Date 头")
    print("=" * 80)

    # 从环境变量读取配置
    access_key = os.getenv('HIAGENT_ACCESS_KEY')
    secret_key = os.getenv('HIAGENT_SECRET_KEY')
    host = os.getenv('HIAGENT_HOST', 'open.volcengineapi.com')
    account_id = os.getenv('HIAGENT_ACCOUNT_ID')
    service = 'llmops'
    region = 'cn-beijing'
    version = '2024-05-01'

    # 创建签名参数
    param = SignParam()
    param.path = "/"
    param.method = "POST"
    param.host = host

    # 设置查询参数
    query = OrderedDict()
    query['Action'] = 'ListAppCenterCategory'
    query['Version'] = version
    query['X-Account-Id'] = account_id
    param.query = query

    # 设置请求头（手动添加X-Date）
    headers = OrderedDict()
    headers['Host'] = host
    headers['Content-Type'] = 'application/json'
    headers['X-Date'] = formatdate(timeval=None, localtime=False, usegmt=True)
    param.header_list = headers
    param.headers = headers

    # 设置请求体
    body = {"Top": {}}
    param.body = json.dumps(body)

    print(f"\n[INFO] 签名前的请求头（已包含X-Date）:")
    for key, value in param.headers.items():
        print(f"  {key}: {value}")

    # 执行签名
    print(f"\n[INFO] 执行签名...")
    signer = SignerV4()
    credentials = Credentials(access_key, secret_key, service, region)

    try:
        signer.sign(param, credentials)
        print(f"[SUCCESS] 签名成功")
    except Exception as e:
        print(f"[ERROR] 签名失败: {e}")
        import traceback
        traceback.print_exc()
        return False

    # 检查签名后的请求头
    print(f"\n[INFO] 签名后的请求头:")
    for key, value in param.headers.items():
        print(f"  {key}: {value}")

    # 尝试发送请求
    print(f"\n[INFO] 发送测试请求...")
    url = f"https://{host}/?Action=ListAppCenterCategory&Version={version}&X-Account-Id={account_id}"

    try:
        response = requests.post(
            url,
            headers=dict(param.headers),
            data=param.body,
            timeout=30,
            verify=False
        )

        print(f"[INFO] 响应状态码: {response.status_code}")
        print(f"\n[INFO] 响应内容:")
        try:
            result = response.json()
            print(json.dumps(result, indent=2, ensure_ascii=False))

            if response.status_code == 200:
                print(f"\n[SUCCESS] 请求成功!")
                return True
            else:
                print(f"\n[ERROR] 请求失败!")
                return False

        except json.JSONDecodeError:
            print(response.text)
            return False

    except Exception as e:
        print(f"[ERROR] 请求失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    print("\n" + "=" * 80)
    print("X-Date 头验证测试")
    print("=" * 80)

    # 测试1: 让SDK自动添加X-Date
    success1 = test_sign_with_xdate()

    # 测试2: 手动添加X-Date
    success2 = test_manual_xdate()

    print("\n" + "=" * 80)
    print("测试总结")
    print("=" * 80)
    print(f"自动添加X-Date: {'成功' if success1 else '失败'}")
    print(f"手动添加X-Date: {'成功' if success2 else '失败'}")

    if success2 and not success1:
        print("\n[建议] volcengine SDK 未自动添加 X-Date 头")
        print("        需要在 api_client.py 中手动添加 X-Date 头")
    elif success1:
        print("\n[结论] volcengine SDK 正常工作")
        print("        问题可能出在其他地方")
